import random
import pylab
from PIL import Image
import numpy as np
from matplotlib.pyplot import imshow
import matplotlib.pyplot as plt
import sys
from matplotlib import pyplot
from keras.utils import to_categorical
from keras.models import Sequential
from keras.layers import Conv2D
from keras.layers import MaxPool2D
from keras.layers import Dense
from keras.layers import Flatten
from keras.layers import Dropout
from keras.optimizers import SGD
from keras.preprocessing.image import ImageDataGenerator
from keras.models import load_model
import tensorflow as tf
import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

model_path = "models/vegetable_recognition_cnn.keras"
model = load_model(model_path)

# 设置字体
plt.rcParams['font.sans-serif'] = ['SimHei']

class_names = []
def data_load_class_names(data_dir, img_height, img_width, batch_size):
    train_ds = tf.keras.preprocessing.image_dataset_from_directory(
        data_dir,
        label_mode='categorical',
        seed=123,
        image_size=(img_height, img_width),
        batch_size=batch_size)
    class_names = train_ds.class_names
    # 返回处理之后的训练集、验证集和类名
    return class_names

def read_random_image():
    folder = r'D:\\workspace\\pythonworkspace\\vegetable_recogniton\\data\\validation_image_data\\'
    file_path = folder + random.choice(os.listdir(folder))
    pil_im = Image.open(file_path, 'r')
    return pil_im


def get_predict(img, model):
    img = img.resize((224, 224), Image.BILINEAR)
    img = np.asarray(img)
    outputs = model.predict(img.reshape(1, 224, 224, 3))
    result_index = int(np.argmax(outputs))
    result = class_names[result_index]
    print("结果为", '%s' % result)
    return result

class_names = data_load_class_names("./data/train_image_data", 224, 224, 16)
pil_im = read_random_image()
imshow(np.asarray(pil_im))
plt.title(get_predict(pil_im, model))
pylab.show()

